package com.atguigu.hadoop.examples.mapreduce;


import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.conf.Configured;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.BytesWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableComparable;
import org.apache.hadoop.mapred.ClusterStatus;
import org.apache.hadoop.mapred.JobClient;
import org.apache.hadoop.mapreduce.*;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat;
import org.apache.hadoop.util.Tool;
import org.apache.hadoop.util.ToolRunner;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Random;

public class RandomWriter extends Configured
        implements Tool
{
    public static final String TOTAL_BYTES = "mapreduce.randomwriter.totalbytes";
    public static final String BYTES_PER_MAP = "mapreduce.randomwriter.bytespermap";
    public static final String MAPS_PER_HOST = "mapreduce.randomwriter.mapsperhost";
    public static final String MAX_VALUE = "mapreduce.randomwriter.maxvalue";
    public static final String MIN_VALUE = "mapreduce.randomwriter.minvalue";
    public static final String MIN_KEY = "mapreduce.randomwriter.minkey";
    public static final String MAX_KEY = "mapreduce.randomwriter.maxkey";

    public int run(String[] args)
            throws Exception
    {
        if (args.length == 0) {
            System.out.println("Usage: writer <out-dir>");
            ToolRunner.printGenericCommandUsage(System.out);
            return 2;
        }

        Path outDir = new Path(args[0]);
        Configuration conf = getConf();
        JobClient client = new JobClient(conf);
        ClusterStatus cluster = client.getClusterStatus();
        int numMapsPerHost = conf.getInt("mapreduce.randomwriter.mapsperhost", 10);
        long numBytesToWritePerMap = conf.getLong("mapreduce.randomwriter.bytespermap", 1073741824L);

        if (numBytesToWritePerMap == 0L) {
            System.err.println("Cannot havemapreduce.randomwriter.bytespermap set to 0");
            return -2;
        }
        long totalBytesToWrite = conf.getLong("mapreduce.randomwriter.totalbytes", numMapsPerHost * numBytesToWritePerMap * cluster
                .getTaskTrackers());
        int numMaps = (int)(totalBytesToWrite / numBytesToWritePerMap);
        if ((numMaps == 0) && (totalBytesToWrite > 0L)) {
            numMaps = 1;
            conf.setLong("mapreduce.randomwriter.bytespermap", totalBytesToWrite);
        }
        conf.setInt("mapreduce.job.maps", numMaps);

        Job job = Job.getInstance(conf);

        job.setJarByClass(RandomWriter.class);
        job.setJobName("random-writer");
        FileOutputFormat.setOutputPath(job, outDir);
        job.setOutputKeyClass(BytesWritable.class);
        job.setOutputValueClass(BytesWritable.class);
        job.setInputFormatClass(RandomInputFormat.class);
        job.setMapperClass(RandomMapper.class);
        job.setReducerClass(Reducer.class);
        job.setOutputFormatClass(SequenceFileOutputFormat.class);

        System.out.println("Running " + numMaps + " maps.");

        job.setNumReduceTasks(0);

        Date startTime = new Date();
        System.out.println("Job started: " + startTime);
        int ret = job.waitForCompletion(true) ? 0 : 1;
        Date endTime = new Date();
        System.out.println("Job ended: " + endTime);
        System.out.println("The job took " +
                (endTime
                        .getTime() - startTime.getTime()) / 1000L + " seconds.");

        return ret;
    }

    public static void main(String[] args) throws Exception {
        int res = ToolRunner.run(new Configuration(), new RandomWriter(), args);
        System.exit(res);
    }

    static class RandomMapper extends Mapper<WritableComparable, Writable, BytesWritable, BytesWritable>
    {
        private long numBytesToWrite;
        private int minKeySize;
        private int keySizeRange;
        private int minValueSize;
        private int valueSizeRange;
        private Random random = new Random();
        private BytesWritable randomKey = new BytesWritable();
        private BytesWritable randomValue = new BytesWritable();

        private void randomizeBytes(byte[] data, int offset, int length) {
            for (int i = offset + length - 1; i >= offset; i--)
                data[i] = ((byte)this.random.nextInt(256));
        }

        public void map(WritableComparable key, Writable value, Mapper<WritableComparable, Writable, BytesWritable, BytesWritable>.Context context)
                throws IOException, InterruptedException
        {
            int itemCount = 0;
            while (this.numBytesToWrite > 0L)
            {
                int keyLength = this.minKeySize + (this.keySizeRange != 0 ? this.random
                        .nextInt(this.keySizeRange) :
                        0);
                this.randomKey.setSize(keyLength);
                randomizeBytes(this.randomKey.getBytes(), 0, this.randomKey.getLength());

                int valueLength = this.minValueSize + (this.valueSizeRange != 0 ? this.random
                        .nextInt(this.valueSizeRange) :
                        0);
                this.randomValue.setSize(valueLength);
                randomizeBytes(this.randomValue.getBytes(), 0, this.randomValue.getLength());
                context.write(this.randomKey, this.randomValue);
                this.numBytesToWrite -= keyLength + valueLength;
                context.getCounter(RandomWriter.Counters.BYTES_WRITTEN).increment(keyLength + valueLength);
                context.getCounter(RandomWriter.Counters.RECORDS_WRITTEN).increment(1L);
                itemCount++; if (itemCount % 200 == 0) {
                context.setStatus("wrote record " + itemCount + ". " + this.numBytesToWrite + " bytes left.");
            }
            }

            context.setStatus("done with " + itemCount + " records.");
        }

        public void setup(Mapper<WritableComparable, Writable, BytesWritable, BytesWritable>.Context context)
        {
            Configuration conf = context.getConfiguration();
            this.numBytesToWrite = conf.getLong("mapreduce.randomwriter.bytespermap", 1073741824L);

            this.minKeySize = conf.getInt("mapreduce.randomwriter.minkey", 10);
            this.keySizeRange =
                    (conf
                            .getInt("mapreduce.randomwriter.maxkey", 1000) -
                            this.minKeySize);
            this.minValueSize = conf.getInt("mapreduce.randomwriter.minvalue", 0);
            this.valueSizeRange =
                    (conf
                            .getInt("mapreduce.randomwriter.maxvalue", 20000) -
                            this.minValueSize);
        }
    }

    static class RandomInputFormat extends InputFormat<Text, Text>
    {
        public List<InputSplit> getSplits(JobContext job)
                throws IOException
        {
            List result = new ArrayList();
            Path outDir = FileOutputFormat.getOutputPath(job);

            int numSplits = job
                    .getConfiguration().getInt("mapreduce.job.maps", 1);
            for (int i = 0; i < numSplits; i++) {
                result.add(new FileSplit(new Path(outDir, "dummy-split-" + i), 0L, 1L, (String[])null));
            }

            return result;
        }

        public RecordReader<Text, Text> createRecordReader(InputSplit split, TaskAttemptContext context)
                throws IOException, InterruptedException
        {
            return new RandomRecordReader(((FileSplit)split).getPath());
        }

        static class RandomRecordReader extends RecordReader<Text, Text>
        {
            Path name;
            Text key = null;
            Text value = new Text();

            public RandomRecordReader(Path p) { this.name = p; }


            public void initialize(InputSplit split, TaskAttemptContext context)
                    throws IOException, InterruptedException
            {
            }

            public boolean nextKeyValue()
            {
                if (this.name != null) {
                    this.key = new Text();
                    this.key.set(this.name.getName());
                    this.name = null;
                    return true;
                }
                return false;
            }

            public Text getCurrentKey() {
                return this.key;
            }

            public Text getCurrentValue() {
                return this.value;
            }
            public void close() {
            }

            public float getProgress() {
                return 0.0F;
            }
        }
    }

    static enum Counters
    {
        RECORDS_WRITTEN, BYTES_WRITTEN;
    }
}